#include <dirent.h>
#include <sys/stat.h>
#include <iostream>
#include <fstream>
#include <string>
#include <cstring>
#include <vector>
#include <thread>
#include <mutex>
#include <cmath>
#include <queue>
#include "hsv.h"

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

#include <windows.h>


// 使用 Windows API 创建目录
bool createDirectory(const std::string &path) {
    return CreateDirectoryA(path.c_str(), NULL) || GetLastError() == ERROR_ALREADY_EXISTS;
}

// RGB转HSV，输出v（亮度）范围0-1
float rgbToV(unsigned char r, unsigned char g, unsigned char b) {
    float fr = r / 255.0f;
    float fg = g / 255.0f;
    float fb = b / 255.0f;
    float maxc = std::max(fr, std::max(fg, fb));
    //float minc = std::min(fr, std::min(fg, fb));
    float v = maxc; // 亮度V就是最大值
    return v;
}

void processBinaryImage(std::vector<unsigned char>& outputData, int width, int height) {
    // 方向数组：四连通方向（上、下、左、右）
    const int dx[] = {-1, 1, 0, 0};
    const int dy[] = {0, 0, -1, 1};

    // 创建访问标记矩阵
    std::vector<bool> visited(width * height, false);

    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            int idx = y * width + x;

            // 跳过已访问或非零像素
            if (outputData[idx] != 0 || visited[idx]) continue;

            // 洪水填充准备
            std::queue<std::pair<int, int>> q;
            std::vector<std::pair<int, int>> cluster;
            q.push({x, y});
            visited[idx] = true;
            int count = 0;

            // BFS遍历连通区域
            while (!q.empty()) {
                auto [cx, cy] = q.front();
                q.pop();
                cluster.emplace_back(cx, cy);
                ++count;

                // 遍历四个方向
                for (int i = 0; i < 4; ++i) {
                    int nx = cx + dx[i];
                    int ny = cy + dy[i];
                    int nidx = ny * width + nx;

                    // 边界检查
                    if (nx >= 0 && nx < width && ny >= 0 && ny < height) {
                        if (!visited[nidx] && outputData[nidx] == 0) {
                            visited[nidx] = true;
                            q.push({nx, ny});
                        }
                    }
                }
            }

            // 修改小面积区域
            if (count < 3) {
                for (auto& [px, py] : cluster) {
                    outputData[py * width + px] = 255; // 设为非零值
                }
            }
        }
    }
}

void processImage(const std::string &inputPath, const std::string &outputPath) {
    int width, height, channels;
    unsigned char *data = stbi_load(inputPath.c_str(), &width, &height, &channels, 3);
    if (!data) {
        std::cerr << "Failed to load image: " << inputPath << std::endl;
        return;
    }

    // 输出图像数据，单通道灰度
    std::vector<unsigned char> outputData(width * height);

    for (int i = 0; i < width * height; ++i) {
        unsigned char r = data[i * 3 + 0];
        unsigned char g = data[i * 3 + 1];
        unsigned char b = data[i * 3 + 2];

        //float v = rgbToV(r, g, b) * 255.0f; // v的0-255范围亮度

        auto hsv = rgbToHsl(r, g, b);


        //for gigi 33.0f
        if (hsv.l < 50.0f) {
            outputData[i] = 0;   // 黑色
        } else {
            outputData[i] = 255; // 白色
        }
    }

    //processBinaryImage(outputData, width, height);

    stbi_image_free(data);

    std::string outputFilePath = outputPath + "/" + inputPath.substr(inputPath.find_last_of('/') + 1);

    // 保存为灰度PNG图像
    if (!stbi_write_png(outputFilePath.c_str(), width, height, 1, outputData.data(), width)) {
        std::cerr << "Failed to write output image: " << outputFilePath << std::endl;
        return;
    }

    //std::cout << "Saved processed binary image to: " << outputFilePath << std::endl;
}

// 遍历目录，获取所有 PNG 文件路径
void traverseDirectory(const std::string &inputDir, std::vector<std::string> &imagePaths) {
    DIR *dir = opendir(inputDir.c_str());
    if (!dir) {
        std::cerr << "Failed to open directory: " << inputDir << std::endl;
        return;
    }

    struct dirent *entry;
    while ((entry = readdir(dir)) != nullptr) {
        std::string name = entry->d_name;

        // 忽略当前目录 "." 和父目录 ".."
        if (name == "." || name == "..") {
            continue;
        }

        std::string fullPath = inputDir + "/" + name;

        // 检查文件类型
        struct stat statbuf;
        if (stat(fullPath.c_str(), &statbuf) == 0) {
            if (S_ISREG(statbuf.st_mode)) { // 如果是普通文件
                if (name.size() >= 4 && name.substr(name.size() - 4) == ".png") {
                    imagePaths.push_back(fullPath);
                }
            } else if (S_ISDIR(statbuf.st_mode)) { // 如果是目录
                std::cout << "Skipping directory: " << fullPath << std::endl;
            }
        } else {
            std::cerr << "Failed to stat file: " << fullPath << std::endl;
        }
    }

    closedir(dir);
}

// 每个线程的任务
void processImagesBatch(const std::vector<std::string> &imagePaths, const std::string &outputPath) {
    int icount = 0;
    for (const auto &imagePath : imagePaths) {
        processImage(imagePath, outputPath);
        icount += 1;
        if (icount % 1000 == 0)
        {
            //exit(0);
            std::cout << std::this_thread::get_id()<< " processed:" << icount << std::endl;
        }
    }
}

int main(int argc, char *argv[]) {
    if (argc != 3) {
        std::cerr << "Usage: " << argv[0] << " <input_path> <output_path>" << std::endl;
        return -1;
    }

    std::string inputPath = argv[1];
    std::string outputPath = argv[2];

    // 检查输出目录是否存在，如果不存在则创建
    struct stat statbuf;
    if (stat(outputPath.c_str(), &statbuf) != 0) {
        if (!createDirectory(outputPath)) {
            std::cerr << "Failed to create output directory: " << outputPath << std::endl;
            return -1;
        }
    }

    // 获取所有 PNG 文件路径
    std::vector<std::string> imagePaths;
    traverseDirectory(inputPath, imagePaths);

    if (imagePaths.empty()) {
        std::cerr << "No PNG images found in input directory: " << inputPath << std::endl;
        return -1;
    }

    std::cout << "find image count:" << imagePaths.size() << std::endl;

    // 启动多线程处理
    const int threadCount = 16;
    std::vector<std::thread> threads;
    std::vector<std::vector<std::string>> threadImageLists(threadCount);

    // 平均分配任务到各线程
    for (size_t i = 0; i < imagePaths.size(); ++i) {
        threadImageLists[i % threadCount].push_back(imagePaths[i]);
    }

    // 启动线程
    for (int i = 0; i < threadCount; ++i) {
        threads.emplace_back(processImagesBatch, threadImageLists[i], outputPath);
    }

    // 等待所有线程完成
    for (auto &t : threads) {
        if (t.joinable()) {
            t.join();
        }
    }

    std::cout << "Processing completed." << std::endl;
    return 0;
}
